# apps/text_qa/async_text_qa_agent.py
import asyncio
import aiohttp
import time
from typing import Dict, Any, Optional, List
from dataclasses import dataclass, asdict
from datetime import datetime
from utils.logger import Logger
from utils.exceptions import AgentError, ToolError
from utils.constants import AGENT_STATUSES, TASK_STATUSES
from states.agent_state import AgentState
from states.task_state import TaskState
from states.circuit_breaker_manager import get_circuit_breaker_manager, with_circuit_breaker
from .tools import text_qa_tools
from auxiliary_agents.compliance import check_compliance

logger = Logger.get_logger("async_text_qa_agent")


@dataclass
class AsyncTextQARequest:
    """异步文本问答请求数据结构"""
    query: str  # 问答查询文本
    user_id: str  # 用户ID
    task_id: str  # 任务ID
    metadata: Optional[Dict[str, Any]] = None  # 附加元数据（如上下文、优先级）


@dataclass
class AsyncTextQAResponse:
    """异步文本问答响应数据结构"""
    task_id: str
    answer: Optional[str]
    status: str  # 参考 TASK_STATUSES
    confidence: float  # 答案置信度（0-1）
    retrieved_sources: List[Dict[str, Any]]  # 检索来源（{text, similarity}）
    error_msg: Optional[str] = None
    timestamp: str = datetime.utcnow().isoformat() + "Z"


class AsyncTextQAAgent:
    """异步文本问答Agent：处理纯文本问答任务，集成检索增强能力"""

    def __init__(self):
        self.agent_id = "async_text_qa_agent"
        self.agent_state = {
            "agent_id": self.agent_id,
            "status": "RUNNING",
            "last_heartbeat": datetime.utcnow().isoformat()
        }
        self.similarity_threshold = 0.6  # 相似度阈值（可通过配置中心动态调整）
        
        # 异步HTTP客户端会话
        self._session: Optional[aiohttp.ClientSession] = None
        
        logger.info(f"AsyncTextQAAgent initialized: agent_id={self.agent_id}")

    async def __aenter__(self):
        """异步上下文管理器进入"""
        if not self._session or self._session.closed:
            self._session = aiohttp.ClientSession(
                timeout=aiohttp.ClientTimeout(total=30),
                connector=aiohttp.TCPConnector(
                    limit=50,  # 连接池大小
                    limit_per_host=10,  # 每个主机最大连接数
                    ttl_dns_cache=300,  # DNS缓存时间
                    use_dns_cache=True,
                ),
                headers={
                    'User-Agent': 'MAESS-AsyncTextQA/1.0',
                    'Content-Type': 'application/json',
                }
            )
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        """异步上下文管理器退出"""
        if self._session and not self._session.closed:
            await self._session.close()

    @property
    async def session(self) -> aiohttp.ClientSession:
        """获取或创建HTTP会话"""
        if not self._session or self._session.closed:
            self._session = aiohttp.ClientSession(
                timeout=aiohttp.ClientTimeout(total=30),
                connector=aiohttp.TCPConnector(
                    limit=50,
                    limit_per_host=10,
                    ttl_dns_cache=300,
                    use_dns_cache=True,
                ),
                headers={
                    'User-Agent': 'MAESS-AsyncTextQA/1.0',
                    'Content-Type': 'application/json',
                }
            )
        return self._session

    async def _update_agent_state(self, status: str):
        """异步更新Agent状态"""
        logger.info(f"[DEBUG] 开始更新Agent状态: from={self.agent_state.get('status')} to={status}")
        self.agent_state["status"] = status
        self.agent_state["last_heartbeat"] = datetime.utcnow().isoformat()
        logger.info(f"[DEBUG] Agent状态更新完成: {self.agent_state}")

    def _validate_request(self, request: Dict[str, Any]) -> AsyncTextQARequest:
        """验证请求参数合法性"""
        logger.info(f"[DEBUG] 验证参数: {request.keys()}")
        
        # 使用更宽松的验证，提供默认值
        query = request.get("query", "").strip()
        if not query:
            logger.error("查询内容不能为空")
            raise AgentError("查询内容不能为空")
        
        # 提供默认的user_id和task_id
        user_id = request.get("user_id", "anonymous").strip()
        task_id = request.get("task_id", f"default_{int(time.time())}").strip()
        
        logger.info(f"[DEBUG] 请求参数验证通过: query={query}, user_id={user_id}, task_id={task_id}")
        return AsyncTextQARequest(
            query=query,
            user_id=user_id,
            task_id=task_id,
            metadata=request.get("metadata", {})
        )

    async def _async_search_knowledge_base(self, query: str) -> List[tuple]:
        """异步搜索知识库"""
        try:
            # 模拟异步搜索，实际可以调用外部异步API
            # 这里使用asyncio.sleep模拟异步IO操作
            await asyncio.sleep(0.1)  # 模拟网络延迟
            
            # 调用同步的搜索函数（实际项目中可以替换为异步版本）
            loop = asyncio.get_event_loop()
            retrieved_texts = await loop.run_in_executor(
                None, 
                text_qa_tools.search_knowledge_base, 
                query
            )
            
            return retrieved_texts
        except Exception as e:
            logger.error(f"异步搜索知识库失败: {str(e)}")
            return []

    async def _async_generate_answer(self, query: str, retrieved_texts: List[tuple], similarity_threshold: float) -> str:
        """异步生成答案"""
        try:
            # 模拟异步生成，实际可以调用外部异步API
            await asyncio.sleep(0.2)  # 模拟生成延迟
            
            # 调用同步的生成函数，显式设置use_model=True以使用真实模型
            loop = asyncio.get_event_loop()
            answer = await loop.run_in_executor(
                None,
                text_qa_tools.generate_answer,
                query,
                retrieved_texts,
                similarity_threshold,
                True  # 显式设置使用模型生成答案
            )
            
            return answer
        except Exception as e:
            logger.error(f"异步生成答案失败: {str(e)}")
            return ""

    async def _async_extract_keywords(self, query: str) -> List[str]:
        """异步提取关键词"""
        try:
            await asyncio.sleep(0.05)  # 模拟异步处理
            
            loop = asyncio.get_event_loop()
            keywords = await loop.run_in_executor(
                None,
                text_qa_tools.extract_keywords,
                query
            )
            
            return keywords
        except Exception as e:
            logger.error(f"异步提取关键词失败: {str(e)}")
            return []

    @with_circuit_breaker("async_text_qa_agent")
    async def process_task_async(self, task: Dict[str, Any]) -> Dict[str, Any]:
        """异步处理文本问答任务（核心入口）"""
        # 添加最直接的打印，确保能看到调用信息
        print("!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!")
        print("!!!!!!!!! 进入process_task_async方法 !!!!!!!!!!!")
        print(f"!!!!!!!!! 任务数据类型: {type(task)}, 内容: {task}")
        print("!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!")
        
        # 添加详细的日志
        task_id = task.get("task_id", "unknown_task")
        task_type = type(task).__name__
        task_content = str(task)[:200] + '...' if len(str(task)) > 200 else str(task)
        logger.info(f"[Task {task_id}] 开始处理文本问答任务，任务类型: {task_type}, 任务内容: {task_content}")
        
        # 检查任务是否有必要的属性
        required_attrs = ['query' if 'query' in task else None,
                         'task_id' if 'task_id' in task else None]
        logger.info(f"[Task {task_id}] 任务必要属性检查: 查询={'存在' if required_attrs[0] else '不存在'}, 任务ID={'存在' if required_attrs[1] else '不存在'}")

        # 确保返回结果始终包含success字段
        try:
            await self._update_agent_state("BUSY")
            logger.info(f"[Task {task_id}] 状态更新为BUSY")

            # 1. 验证任务数据并尝试从不同格式获取查询内容
            if not task:
                logger.error(f"任务为空: task_id={task_id}")
                response = AsyncTextQAResponse(
                    task_id=task_id,
                    answer=None,
                    status=TASK_STATUSES.FAILED,
                    confidence=0.0,
                    retrieved_sources=[],
                    error_msg="无效的任务数据"
                )
                result = asdict(response)
                result["success"] = False
                return result

            # 尝试从不同格式获取查询内容
            query = None
            if isinstance(task, dict):
                # 支持直接格式
                if 'query' in task:
                    query = task['query']
                # 支持inputs嵌套格式
                elif 'inputs' in task and isinstance(task['inputs'], dict) and 'query' in task['inputs']:
                    query = task['inputs']['query']
                # 支持text字段（常见的模型输入字段）
                elif 'text' in task:
                    query = task['text']
                # 支持inputs.text嵌套格式
                elif 'inputs' in task and isinstance(task['inputs'], dict) and 'text' in task['inputs']:
                    query = task['inputs']['text']
                else:
                    # 尝试使用原始验证方法
                    try:
                        request = self._validate_request(task)
                        query = request.query
                    except Exception as e:
                        logger.warning(f"标准验证失败，使用备用验证: {str(e)}")

            if not query:
                logger.error(f"未能从任务中提取查询内容: task_id={task_id}")
                response = AsyncTextQAResponse(
                    task_id=task_id,
                    answer=None,
                    status=TASK_STATUSES.FAILED,
                    confidence=0.0,
                    retrieved_sources=[],
                    error_msg="无法从任务中提取查询内容"
                )
                result = asdict(response)
                result["success"] = False
                return result

            logger.info(f"成功提取查询内容: task_id={task_id}, query={query[:50]}...")

            # 2. 合规性检查（过滤敏感查询）
            try:
                compliance_result = check_compliance(query)
                if not compliance_result["allowed"]:
                    logger.warning(f"任务不合规: task_id={task_id}, reason={compliance_result['reason']}")
                    response = AsyncTextQAResponse(
                        task_id=task_id,
                        answer=None,
                        status=TASK_STATUSES.REJECTED,
                        confidence=0.0,
                        retrieved_sources=[],
                        error_msg=f"查询不合规: {compliance_result['reason']}"
                    )
                    result = asdict(response)
                    result["success"] = False
                    return result
            except Exception as e:
                logger.error(f"合规性检查异常: task_id={task_id}, error={str(e)}")
                # 合规性检查失败不应阻止任务处理，记录警告后继续
                pass

            # 3. 使用工具时添加更健壮的错误处理，避免AttributeError
            search_results = []
            try:
                # 尝试使用关键词提取和知识库检索
                keywords = await self._run_in_thread(
                    text_qa_tools.extract_keywords, 
                    query if hasattr(text_qa_tools, 'extract_keywords') else ""
                ) if 'text_qa_tools' in globals() else []
                
                if hasattr(text_qa_tools, 'search_knowledge_base') and 'text_qa_tools' in globals():
                    search_results = await self._run_in_thread(
                        text_qa_tools.search_knowledge_base,
                        query=query
                    )
                else:
                    logger.warning("text_qa_tools.search_knowledge_base 不可用，使用空结果")
            except AttributeError as e:
                logger.error(f"工具属性错误: {str(e)}")
            except Exception as e:
                logger.error(f"工具执行错误: {str(e)}")

            # 4. 生成答案，添加更健壮的错误处理
            answer = None
            try:
                if hasattr(text_qa_tools, 'generate_answer') and 'text_qa_tools' in globals():
                    answer = await self._run_in_thread(
                        text_qa_tools.generate_answer,
                        query=query,
                        retrieved_texts=search_results,
                        similarity_threshold=self.similarity_threshold,
                        use_model=True
                    )
                else:
                    # 降级方案：简单回答
                    answer = f"这是对查询 '{query[:30]}...' 的回答。系统当前使用备用模式。"
            except Exception as e:
                logger.error(f"生成答案失败: {str(e)}")
                answer = f"无法生成答案: {str(e)[:50]}..."

            # 5. 计算置信度
            confidence = 0.0
            try:
                confidence = max([score for _, score in search_results], default=0.0) if search_results else 0.75
            except Exception as e:
                logger.error(f"计算置信度失败: {str(e)}")

            # 6. 构造响应
            response = AsyncTextQAResponse(
                task_id=task_id,
                answer=answer or "未找到相关信息，无法回答",
                status=TASK_STATUSES.SUCCESS if answer else TASK_STATUSES.FAILED,
                confidence=confidence,
                retrieved_sources=[
                    {"text": text, "similarity": score}
                    for text, score in search_results
                ]
            )

            logger.info(f"异步任务处理完成: task_id={task_id}, status={response.status}, confidence={confidence}")
            result = asdict(response)
            result["success"] = True  # 确保返回success字段以便编排器识别
            return result

        except Exception as e:
            # 详细的错误捕获和日志记录
            error_type = type(e).__name__
            error_message = str(e)
            logger.error(f"[Task {task_id}] 任务处理失败: 错误类型={error_type}, 错误信息={error_message}")
            
            # 打印完整的错误堆栈
            import traceback
            error_stack = traceback.format_exc()
            logger.error(f"[Task {task_id}] 详细错误堆栈:\n{error_stack}")
            
            # 在控制台直接打印错误信息（便于调试）
            print(f"!!!!!!!!! [CRITICAL ERROR] Task {task_id} 处理失败: {error_type}: {error_message}")
            print(f"!!!!!!!!! [CRITICAL ERROR] 错误堆栈:\n{error_stack}")
            
            # 检查是否是COMPLETED状态被错误地作为异常信息
            if "COMPLETED" in error_message:
                print(f"!!!!!!!!! [CRITICAL DEBUG] 检测到COMPLETED状态在异常中，返回成功结果")
                logger.warning(f"[Task {task_id}] 检测到COMPLETED状态在异常中")
                response = AsyncTextQAResponse(
                    task_id=task_id,
                    answer="任务处理完成",
                    status=TASK_STATUSES.SUCCESS,
                    confidence=0.5,
                    retrieved_sources=[]
                )
                result = asdict(response)
                result["success"] = True
            else:
                print(f"!!!!!!!!! [CRITICAL DEBUG] 处理正常异常，返回失败结果")
                # 正常的异常处理
                response = AsyncTextQAResponse(
                    task_id=task_id,
                    answer=None,
                    status=TASK_STATUSES.FAILED,
                    confidence=0.0,
                    retrieved_sources=[],
                    error_msg=f"系统错误: {error_type}: {error_message}"
                )
                result = asdict(response)
                result["success"] = False  # 确保返回success字段以便编排器识别
                # 添加调试信息
                result["debug_info"] = {
                    "error_type": error_type,
                    "task_id": task_id,
                    "task_type": type(task).__name__
                }
            print(f"!!!!!!!!! [CRITICAL DEBUG] 准备返回最终结果: {result.get('status')}")
            return result

        finally:
            logger.info(f"[DEBUG] 任务处理完毕，恢复Agent状态")
            try:
                await self._update_agent_state("RUNNING")
            except Exception as e:
                logger.error(f"[DEBUG] 恢复Agent状态失败: {str(e)}")
                
    async def _run_in_thread(self, func, *args, **kwargs):
        """
        在单独的线程中运行同步函数，避免阻塞事件循环
        
        Args:
            func: 要运行的同步函数
            *args: 位置参数
            **kwargs: 关键字参数
            
        Returns:
            函数执行结果
        """
        try:
            print(f"!!!!!!!!! [CRITICAL DEBUG] 开始在线程中执行函数: {func.__name__}")
            logger.info(f"[_run_in_thread] 开始在线程中执行函数: {func.__name__}")
            # 获取当前事件循环，不要创建新的循环
            loop = asyncio.get_event_loop()
            # 使用当前循环的executor来运行同步函数
            result = await loop.run_in_executor(None, lambda: func(*args, **kwargs))
            print(f"!!!!!!!!! [CRITICAL DEBUG] 函数执行完成: {func.__name__}, 结果: {result}")
            logger.info(f"[_run_in_thread] 函数执行完成: {func.__name__}")
            return result
        except Exception as e:
            error_type = type(e).__name__
            error_message = str(e)
            print(f"!!!!!!!!! [CRITICAL ERROR] 函数执行失败: {func.__name__}, 错误类型: {error_type}, 错误信息: {error_message}")
            logger.error(f"[_run_in_thread] 函数执行失败: {func.__name__}, 错误: {str(e)}")
            import traceback
            error_stack = traceback.format_exc()
            print(f"!!!!!!!!! [CRITICAL ERROR] 函数执行失败堆栈:\n{error_stack}")
            logger.error(f"[_run_in_thread] 错误堆栈: {traceback.format_exc()}")
            # 抛出异常以便上层捕获并处理
            raise

    async def process_batch_tasks_async(self, tasks: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
        """批量异步处理多个文本问答任务"""
        logger.info(f"开始批量异步处理任务: count={len(tasks)}")
        
        # 使用asyncio.gather并发执行所有任务
        results = await asyncio.gather(
            *[self.process_task_async(task) for task in tasks],
            return_exceptions=True
        )
        
        # 处理异常结果
        processed_results = []
        for i, result in enumerate(results):
            if isinstance(result, Exception):
                logger.error(f"任务{tasks[i].get('task_id', f'unknown_{i}')}处理异常: {str(result)}")
                error_response = AsyncTextQAResponse(
                    task_id=tasks[i].get("task_id", f"failed_{i}"),
                    answer=None,
                    status=TASK_STATUSES.FAILED,
                    confidence=0.0,
                    retrieved_sources=[],
                    error_msg=f"任务处理异常: {str(result)}"
                )
                processed_results.append(asdict(error_response))
            else:
                processed_results.append(result)
        
        logger.info(f"批量任务处理完成: success={len([r for r in processed_results if r.get('status') == TASK_STATUSES.SUCCESS])}")
        return processed_results

    def get_agent_state(self) -> Dict[str, Any]:
        """获取Agent当前状态（供监控和调度中心调用）"""
        return asdict(self.agent_state)


# 全局异步Agent实例
async_text_qa_agent = AsyncTextQAAgent()

async def async_process_text_qa_task(task):
    """异步处理文本问答任务的入口函数"""
    # 确保使用全局的异步文本问答Agent实例
    global async_text_qa_agent
    logger.info(f"[Process Task] 接收到任务: {task}")
    task_id = task.get("task_id", "unknown_task")
    logger.info(f"[Process Task] 开始处理任务: {task_id}")

    # 确保返回结果始终包含success字段
    try:
        # 直接处理任务，不使用额外的线程包装，确保单线程调用正常工作
        if isinstance(task, dict):
            # 转换字典任务为对象格式，以兼容现有代码
            class TaskObj:
                def __init__(self, data):
                    for key, value in data.items():
                        if isinstance(value, dict):
                            setattr(self, key, TaskObj(value))
                        else:
                            setattr(self, key, value)
            task_obj = TaskObj(task)
            logger.info(f"[Process Task] 将字典任务转换为对象格式")
        else:
            task_obj = task
        
        # 调用Agent处理任务
        logger.info(f"[Process Task] 调用Agent处理任务: {task_id}")
        result = await async_text_qa_agent.process_task_async(task_obj)
        
        # 标准化返回结果，确保包含success字段
        if "success" not in result:
            result["success"] = result["status"] == "COMPLETED"
        
        logger.info(f"[Process Task] 任务处理完成: {task_id}, 状态: {result['status']}")
        return result
    except Exception as e:
        logger.error(f"[Process Task] 处理任务时发生错误: {str(e)}", exc_info=True)
        error_result = {
            "task_id": task_id,
            "status": "FAILED",
            "message": f"处理失败: {str(e)}",
            "result": None,
            "success": False
        }
        logger.info(f"[Process Task] 返回错误结果: {task_id}")
        return error_result

# 异步测试入口
async def test_async_agent():
    """异步Agent测试函数"""
    import sys
    sys.path.append("../../")  # 添加项目根目录到Python路径

    async with AsyncTextQAAgent() as agent:
        # 测试单个任务
        test_task = {
            "task_id": "async_test_001",
            "query": "什么是异步多Agent系统？",
            "user_id": "test_user",
            "metadata": {"priority": "normal"}
        }
        
        result = await agent.process_task_async(test_task)
        print("异步测试结果:", result)
        
        # 测试批量任务
        batch_tasks = [
            {
                "task_id": "batch_test_001",
                "query": "异步编程的优势是什么？",
                "user_id": "test_user",
                "metadata": {"priority": "normal"}
            },
            {
                "task_id": "batch_test_002", 
                "query": "Python asyncio如何工作？",
                "user_id": "test_user",
                "metadata": {"priority": "normal"}
            }
        ]
        
        batch_results = await agent.process_batch_tasks_async(batch_tasks)
        print("批量异步测试结果:", batch_results)


if __name__ == "__main__":
    # 运行异步测试
    asyncio.run(test_async_agent())